#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>

#include <cuda.h>
#include <cuda_runtime.h>

#include <vector>

namespace {
    __global__ void check_win_cuda_kernel(
    torch::PackedTensorAccessor<bool,4,torch::RestrictPtrTraits,size_t> board,
    torch::PackedTensorAccessor<int64_t,1,torch::RestrictPtrTraits,size_t> action,
    torch::PackedTensorAccessor<bool,1,torch::RestrictPtrTraits,size_t> win){
    const static int dirs[4][2] = {{1, 0}, {0, 1}, {1, 1}, {1, -1}};
    const int i = blockIdx.x * blockDim.x + threadIdx.x;

    auto w = board.size(3);
    auto j = action[i] / w;
    auto k = action[i] % w;
    board[i][0][j][k] = 1;
    win[i] = 0;
    for (auto& dir : dirs) {
      int count = 1;
      int x = j + dir[0], y = k + dir[1];
      while (x >= 0 && y >= 0 && x < board.size(2) && y < board.size(3) &&
             board[i][0][x][y])
        x += dir[0], y += dir[1], count++;
      x = j - dir[0], y = k - dir[1];
      while (x >= 0 && y >= 0 && x < board.size(2) && y < board.size(3) &&
             board[i][0][x][y])
        x -= dir[0], y -= dir[1], count++;
      if (count >= 5) {
        win[i] = 1;
        break;
      }
    }
}
}

torch::Tensor check_win_cuda(torch::Tensor board, torch::Tensor action){
    const int threads = min(1024, (int)board.size(0));
    const dim3 blocks((board.size(0) + threads - 1) / threads);
    torch::Tensor win = torch::empty_like(action, torch::TensorOptions().dtype(torch::kBool));
    auto stream = at::cuda::getCurrentCUDAStream();
    check_win_cuda_kernel<<<blocks, threads, 0, stream>>>(
        board.packed_accessor<bool,4,torch::RestrictPtrTraits,size_t>(),
        action.packed_accessor<int64_t,1,torch::RestrictPtrTraits,size_t>(),
        win.packed_accessor<bool,1,torch::RestrictPtrTraits,size_t>());
    return win;
}
